package com.leetcode.partition3;

import java.util.HashMap;
import java.util.Map;

/**
 * @author `RKC`
 * @date 2021/9/13 18:13
 */
public class LC208实现Trie {

    private static class Trie {

        private TrieNode root;

        /** Initialize your data structure here. */
        public Trie() {
            root = new TrieNode(' ');
        }

        /** Inserts a word into the trie. */
        public void insert(String word) {
            char[] words = word.toCharArray();
            TrieNode cur = root;
            for (int i = 0; i < words.length; i++) {
                if (cur.trieNodeMap.containsKey(words[i])) {
                    if (i == words.length - 1) cur.trieNodeMap.get(words[i]).end = true;
                    cur = cur.trieNodeMap.get(words[i]);
                    continue;
                }
                TrieNode trieNode = new TrieNode(words[i]);
                if (i == words.length - 1) trieNode.end = true;
                cur.trieNodeMap.put(words[i], trieNode);
                cur = cur.trieNodeMap.get(words[i]);
            }
        }

        /** Returns if the word is in the trie. */
        public boolean search(String word) {
            TrieNode cur = root;
            for (char ch : word.toCharArray()) {
                if (cur.trieNodeMap.containsKey(ch)) cur = cur.trieNodeMap.get(ch);
                else return false;
            }
            return cur.end;
        }

        /** Returns if there is any word in the trie that starts with the given prefix. */
        public boolean startsWith(String prefix) {
            TrieNode cur = root;
            for (char ch : prefix.toCharArray()) {
                if (cur.trieNodeMap.containsKey(ch)) cur = cur.trieNodeMap.get(ch);
                else return false;
            }
            return true;
        }
    }

    private static class TrieNode {
        char data;
        Map<Character, TrieNode> trieNodeMap;
        boolean end;

        public TrieNode(char data) {
            this.data = data;
            trieNodeMap = new HashMap<>();
            end = false;
        }
    }
}
